package com.aloudmeta.grpc.support.metric;

import io.grpc.*;
import io.grpc.Status.Code;
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Timer;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.protocol.grpc.interceptors.ServerInterceptor;

import java.util.EnumMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;

import static com.aloudmeta.grpc.support.metric.MetricConstants.*;
import static org.apache.dubbo.common.constants.CommonConstants.PROVIDER;

/**
 * huaixin 2022/2/9 2:00 PM
 */
@Slf4j
@Activate(group = {PROVIDER}, order=2500)
public class MetricCollectingServerInterceptor implements ServerInterceptor {

    private final Map<MethodDescriptor<?, ?>, MetricSet> metricsForMethods = new ConcurrentHashMap<>();

    private MeterRegistry simpleMeterRegistry;

    private BindableService dubboEchoService;

    public void setSimpleMeterRegistry(MeterRegistry simpleMeterRegistry) {
        this.simpleMeterRegistry = simpleMeterRegistry;
    }

    public void setDubboEchoService(BindableService dubboEchoService) {
        preregisterService(dubboEchoService.bindService());
    }

    public void preregisterService(final ServerServiceDefinition serviceDefinition) {
        preregisterService(serviceDefinition.getServiceDescriptor());
    }

    public void preregisterService(final ServiceDescriptor service) {
        for (final MethodDescriptor<?, ?> method : service.getMethods()) {
            preregisterMethod(method);
        }
    }

    public void preregisterMethod(final MethodDescriptor<?, ?> method) {
        metricsFor(method);
    }

    protected Counter newRequestCounterFor(final MethodDescriptor<?, ?> method) {
        return MetricUtils.prepareCounterFor(method,
                METRIC_NAME_SERVER_REQUESTS_RECEIVED,
                "The total number of requests received").register(this.simpleMeterRegistry);
    }

    protected Counter newResponseCounterFor(final MethodDescriptor<?, ?> method) {
        return MetricUtils.prepareCounterFor(method,
                METRIC_NAME_SERVER_RESPONSES_SENT,
                "The total number of responses sent").register(this.simpleMeterRegistry);
    }

    protected Function<Code, Timer> newTimerFunction(final MethodDescriptor<?, ?> method) {
        return asTimerFunction(() -> MetricUtils.prepareTimerFor(method,
                METRIC_NAME_SERVER_PROCESSING_DURATION,
                "The total time taken for the server to complete the call"));
    }

    @Override
    public <Q, A> ServerCall.Listener<Q> interceptCall(
            final ServerCall<Q, A> call,
            final Metadata requestHeaders,
            final ServerCallHandler<Q, A> next) {

        final MetricSet metrics = metricsFor(call.getMethodDescriptor());
        final Consumer<Code> responseStatusTiming = metrics.newProcessingDurationTiming(this.simpleMeterRegistry);

        final MetricCollectingServerCall<Q, A> monitoringCall =
                new MetricCollectingServerCall<>(call, metrics.getResponseCounter());

        return new MetricCollectingServerCallListener<>(
                next.startCall(monitoringCall, requestHeaders),
                metrics.getRequestCounter(),
                monitoringCall::getResponseCode,
                responseStatusTiming);
    }

    /**
     * Gets or creates a {@link MetricSet} for the given gRPC method. This will initialize all default counters and
     * timers for that method.
     *
     * @param method The method to get the metric set for.
     * @return The metric set for the given method.
     * @see #newMetricsFor(MethodDescriptor)
     */
    protected final MetricSet metricsFor(final MethodDescriptor<?, ?> method) {
        return this.metricsForMethods.computeIfAbsent(method, this::newMetricsFor);
    }

    /**
     * Creates a {@link MetricSet} for the given gRPC method. This will initialize all default counters and timers for
     * that method.
     *
     * @param method The method to get the metric set for.
     * @return The newly created metric set for the given method.
     */
    protected MetricSet newMetricsFor(final MethodDescriptor<?, ?> method) {
        log.debug("Creating new metrics for {}", method.getFullMethodName());
        return new MetricSet(newRequestCounterFor(method), newResponseCounterFor(method), newTimerFunction(method));
    }

    protected Function<Code, Timer> asTimerFunction(final Supplier<Timer.Builder> timerTemplate) {
        final Map<Code, Timer> cache = new EnumMap<>(Code.class);
        final Function<Code, Timer> creator = code -> timerTemplate.get()
                .tag(TAG_STATUS_CODE, code.name())
                .register(this.simpleMeterRegistry);
        final Function<Code, Timer> cacheResolver = code -> cache.computeIfAbsent(code, creator);
        cacheResolver.apply(Code.OK);
        return cacheResolver;
    }

    /**
     * Container for all metrics of a certain call. Used instead of 3 maps to improve performance.
     */
    @Getter
    protected static class MetricSet {

        private final Counter requestCounter;
        private final Counter responseCounter;
        private final Function<Code, Timer> timerFunction;

        /**
         * Creates a new metric set with the given meter instances.
         *
         * @param requestCounter The request counter to use.
         * @param responseCounter The response counter to use.
         * @param timerFunction The timer function to use.
         */
        public MetricSet(
                final Counter requestCounter,
                final Counter responseCounter,
                final Function<Code, Timer> timerFunction) {

            this.requestCounter = requestCounter;
            this.responseCounter = responseCounter;
            this.timerFunction = timerFunction;
        }

        /**
         * Uses the given registry to create a {@link Timer.Sample Timer.Sample} that will be reported if the returned
         * consumer is invoked.
         *
         * @param registry The registry used to create the sample.
         * @return The newly created consumer that will report the processing duration since calling this method and
         *         invoking the returned consumer along with the status code.
         */
        public Consumer<Status.Code> newProcessingDurationTiming(final MeterRegistry registry) {
            final Timer.Sample timerSample = Timer.start(registry);
            return code -> timerSample.stop(this.timerFunction.apply(code));
        }

    }
}
